import numpy as np
import matplotlib.pyplot as plt
import sklearn.linear_model as lm
import sklearn.datasets as ds
import sklearn.model_selection as ms

from utils import utils

bd = ds.load_boston()
# 获取波士顿房价的所有特征数据
data = bd.data
# 获取每行特征对应的房价
label = bd.target
# RM：住宅平均房间数
nox = data[:, 5:6]
# 将数据拆分成80%的训练数据  20%的测试数据
xtrain, xtest, ytrain, ytest = ms.train_test_split(nox, label, test_size=0.2, random_state=10)
# 将 [[4],[2]]这样的特征矩阵转换成 [4,2]这样的向量 绘制散点图
# plt.scatter(xtrain[:, -1], ytrain, c="red")
# plt.show()

xtrain = xtrain[ytrain < 50]
ytrain = ytrain[ytrain < 50]

# 创建线程回归的类
lr = lm.LinearRegression()
lr.fit(xtrain, ytrain)
# 系数也就是斜率
print(lr.coef_)
# 截距
print(lr.intercept_)


plt.scatter(xtrain[:, -1], ytrain, c="red")
# 绘制80%的真实数据
plt.plot(xtrain[:, -1], xtrain[:, -1] * lr.coef_ + lr.intercept_)
plt.show()
plt.close()

my_model = utils.LinearRegression()
my_model.fit(xtrain, ytrain)
ypred = my_model.predict(xtest)

plt.scatter(xtrain, ytrain, marker='o', s=25, edgecolor='k')
plt.plot(xtest, ypred, color='blue', linewidth=3)
plt.show()